import cv2
import tkinter as tk
from tkinter import filedialog, Tk, Button, StringVar, OptionMenu, ttk, messagebox
from tkinter.filedialog import askopenfilename
import imutils
import numpy as np
####202112405020 蔡丹真########
#最终完成版
# 选择视频文件
def choose_video():
    video_file = filedialog.askopenfilename(title="选择视频文件", filetypes=[("MP4文件", "*.mp4")])
    return video_file

# 选择图片文件
def choose_files():
    images_paths = filedialog.askopenfilenames(title="选择图片",
                                               filetypes=[ ("PNG Files", "*.png"),("JPEG Files", "*.jpg")])
    return images_paths

# 从视频中平均提取帧并确保帧之间有重叠
def extract_frames(video_path, num_frames=16):
    cap = cv2.VideoCapture(video_path)
    frames = []

    # 获取视频的总帧数
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if total_frames < num_frames:
        print("视频帧数不足，无法提取指定数量的帧。")
        cap.release()
        return []

    # 计算每帧之间的间隔（均匀分布）
    step = total_frames / num_frames

    for i in range(num_frames):
        # 计算当前帧的索引
        frame_idx = int(i * step)
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)  # 设置视频读取的位置

        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)

    cap.release()
    return frames


# 拼接图像
def stitch_images(images):
    stitcher = cv2.Stitcher_create()
    status, panorama = stitcher.stitch(images)
    return status, panorama


# 选择存储类型并保存全景图
def chose_savetype(panorama):
    def on_save():
        # 获取用户选择的格式
        save_format = save_format_var.get()

        # 弹出保存路径对话框
        save_path = filedialog.asksaveasfilename(defaultextension=f".{save_format}",
                                                 filetypes=[(f"{save_format.upper()} Files", f"*.{save_format}")])
        if save_path:
            # 保存全景图
            cv2.imwrite(save_path, panorama)
            print(f"全景图已保存为: {save_path}")

        save_window.destroy()

    # 创建保存类型选择窗口
    save_window = Tk()
    save_window.title("选择保存格式")

    # 存储格式选择
    save_format_var = StringVar(save_window)
    save_format_var.set("png")  # 默认选择png格式

    format_options = ["png", "jpg", "hdr"]  # 可选的保存格式
    save_format_menu = OptionMenu(save_window, save_format_var, *format_options)
    save_format_menu.pack(pady=20)

    # 保存按钮
    save_button = Button(save_window, text="保存", command=on_save)
    save_button.pack(pady=20)

    save_window.mainloop()


# 主界面，提供选择视频或图片的按钮
def main_window():
    # 创建主窗口
    root = tk.Tk()
    root.title("全景图拼接小工具")

    # 创建一个标签
    label = tk.Label(root, text="欢迎使用，请选择原始素材", font=("Arial", 14), bg="#F7F1E7", fg="#6B4F30")

    label.pack(pady=20)
    # 设置窗口大小
    root.geometry("500x550")
    root.config(bg="#F7F1E7")
    # 使用主题
    style = ttk.Style()
    style.theme_use('alt')
    style.configure('TButton',
                    bfont=("Poppins", 10),  # 字体样式
                    foreground="#6B4F30",  # 前景色（文字颜色）
                    background="#F1D1B2",  # 背景色
                    relief="flat",  # 按钮的边框样式
                    padding=10  # 按钮的内边距
                    )

    def on_video_button_click():
        video_path = choose_video()
        if video_path:
            process_video(video_path)

    def on_image_button_click():
        images_paths = choose_files()
        if len(images_paths) < 2:
            print("请选择至少两张图片来进行拼接！")
            return
        process_images(images_paths)

    def on_smooth_button_click():
        Tk().withdraw()  # 不显示主窗口
        image_path = askopenfilename(title="Select an Image", filetypes=[("Image files", "*.jpg;*.jpeg;*.png")])
        if image_path:
            remove_black_borders(image_path)

    def exit_program():
        root.quit()  # 退出程序

    # 创建两个按钮
    video_button = ttk.Button(root, text="选择视频", command=on_video_button_click, style="TButton")
    video_button.pack(pady=20)

    image_button = ttk.Button(root, text="选择图片", command=on_image_button_click, style="TButton")
    image_button.pack(pady=20)

    smooth_button = ttk.Button(root, text="去除黑边", command=on_smooth_button_click, style="TButton")
    smooth_button.pack(pady=20)
    # 创建退出按钮
    exit_button = ttk.Button(root, text="退出", command=exit_program, style="TButton")
    exit_button.pack(pady=20)

    root.mainloop()


# 处理视频文件
def process_video(video_path):
    # 提取视频帧
    frames = extract_frames(video_path)

    if len(frames) < 2:
        print("视频中的帧太少，无法生成全景图！")
        return

    # 拼接图像
    status, panorama = stitch_images(frames)

    if status == cv2.Stitcher_OK:
        # 拼接成功，显示拼接后的图片
        print("拼接成功！")

        # 显示拼接后的全景图
        cv2.imshow("Panorama_fromVedio", panorama)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

        # 保存全景图
        chose_savetype(panorama)
    else:
        print("拼接失败，请检查视频质量或帧之间的重叠区域。")



# 处理图片文件
def process_images(images_paths):
    images = []
    for image_path in images_paths:
        images.append(cv2.imread(image_path))

    # 检查图片是否成功读取
    for i, img in enumerate(images):
        if img is None:
            print(f"警告: 图片 {images_paths[i]} 无法读取，请检查文件路径或格式")
            return

    # 执行拼接
    status, panorama = stitch_images(images)
    def show_error_message():
        root = tk.Tk()
        root.withdraw()  # 隐藏主窗口
        messagebox.showerror("拼接失败", "拼接失败，请检查图片质量或确保图片之间有足够的重叠区域。")
        root.quit()
    if status == cv2.Stitcher_OK:
        # 拼接成功，显示并调用保存格式选择函数
        print("拼接成功！")

        # 显示拼接后的图片
        cv2.imshow("Panorama", panorama)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

        # 选择保存格式并保存
        chose_savetype(panorama)
    else:
        # 拼接失败
        print("拼接失败，请检查图片质量或确保图片之间有足够的重叠区域。")
        # 拼接失败，弹出提示窗口
        show_error_message()
# 定义去除黑边的函数
def remove_black_borders(image_path):
    stitched_image = cv2.copyMakeBorder(cv2.imread(image_path),10,10,10,10,cv2.BORDER_CONSTANT,(0,0,0))

    # 将图像转为灰度图像
    gray = cv2.cvtColor(stitched_image, cv2.COLOR_BGR2GRAY)
    thresh_img = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)[1]
    # 阈值化，将黑色区域变为白色

    # 找到图像中的轮廓
    contours= cv2.findContours(thresh_img.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = imutils.grab_contours(contours)
    areaOI = max(contours, key=cv2.contourArea)

    mask =np.zeros(thresh_img.shape, dtype="uint8")
    # 获取最大轮廓的矩形边界框
    x, y, w, h = cv2.boundingRect(areaOI)
    cv2.rectangle(mask, (x, y), (x + w, y + h), 255, -1)

    minRectangle = mask.copy()
    sub = mask.copy()

    while cv2.countNonZero(sub) > 0:
        minRectangle = cv2.erode(minRectangle, None)
        sub = cv2.subtract(minRectangle, thresh_img)

    contours  = cv2.findContours(minRectangle.copy(),cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = imutils.grab_contours(contours)
    areaOI = max(contours, key=cv2.contourArea)

    # 显示平滑后的图像
    cv2.imshow("minRectangle image", minRectangle)
    cv2.waitKey(0)
    x, y, w, h = cv2.boundingRect(areaOI)
    stitched_image = stitched_image[y:y+h, x:x+w]
    cv2.imshow("stitched And Processed Image", stitched_image)
    cv2.waitKey(0)
# 保存全景图
    chose_savetype(stitched_image)

# 程序入口
if __name__ == "__main__":
    main_window()
